import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from LSTM_one.LSTM.dataset_create import create_dataset
from LSTM_one.LSTM.lstm import LSTM

plt.rcParams['font.sans-serif']=['SimHei']  #图中字体改为黑体以兼容中文
plt.rcParams['axes.unicode_minus']=False #负号显示的问题
# 读取数据
data = pd.read_csv('../data/2206_2404新能源汽车总体销量数据_全国.csv',encoding='latin1')

# 归一化处理
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(data['sale_count'].values.reshape(-1,1))

# 创建训练集和测试集的有监督学习数据
look_back = 6
train_X, train_Y = create_dataset(scaled_data, look_back)
# 将数据转换为张量
train_X = torch.from_numpy(train_X).type(torch.Tensor)
train_Y = torch.from_numpy(train_Y).type(torch.Tensor)
# 定义模型参数
input_size = 1
hidden_size = 120
num_layers = 3
output_size = 1
learning_rate = 0.001
num_epochs = 800

# 定义模型、损失函数和优化器
lstm = LSTM(input_size, hidden_size, num_layers, output_size)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(lstm.parameters(), lr=learning_rate)

# 训练模型
train_loss = []
for epoch in range(num_epochs):
    outputs = lstm(train_X)
    optimizer.zero_grad()
    loss = criterion(outputs, train_Y)
    loss.backward()
    optimizer.step()
    print('Epoch [{}/{}], Train Loss: {:.4f}'.format(epoch + 1, num_epochs, loss.item()))
    train_loss.append(loss.item())


# 反归一化
train_predict = lstm(train_X)
train_predict = scaler.inverse_transform(train_predict.detach().numpy())
train_predict = np.round(train_predict)
train_Y = scaler.inverse_transform(train_Y.detach().numpy())

# 计算MAPE
def mape(y_true, y_pred):
    mape=[]
    for i in range(len(y_true)):
        mape_=np.abs((y_true[i] - y_pred[i]) / y_true[i])
        mape.append(mape_)
    return np.array(mape),np.mean(mape)

train_real=np.array(data['sale_count'][look_back:])
train_predict=np.array(train_predict)
mape_,mean_mape=mape(train_real,train_predict)
print('平均相对误差:',mean_mape)
#求准确率
def accurecy(y_true,y_pred):
    accuracy=[]
    for i in range(len(y_true)):
        accuracy_=1-np.abs((y_true[i]-y_pred[i])/y_true[i])
        accuracy.append(accuracy_)
    return np.array(accuracy),np.mean(accuracy)

accuracy_,accuracy = accurecy(train_real,train_predict)
print('accuracy:',accuracy)
#模型存储
torch.save(lstm.state_dict(),'../model/lstm_model.pth')
# # 绘制预测结果图
plt.plot(pd.date_range(start='2023-01', end='2024-05', freq='M'),train_real, label='真实值')
plt.plot(pd.date_range(start='2023-01', end='2024-05', freq='M'),train_predict,label='预测值')
plt.xlabel('时间')
plt.ylabel('销量')
plt.title('LSTM模型实际值与预测值比较')
plt.legend()
plt.show()

#存储loss值
file_path = '../loss/loss.txt'
with open(file_path,'w') as file:
    for epoch in range(num_epochs):
        file.write('Epoch [{}/{}], Train Loss: {:.4f}\n'.format(epoch + 1, num_epochs, train_loss[epoch]))
    file.write('平均相对误差:{:.8f}\n'.format(mean_mape))
    file.write('accuracy:{:.8f}\n'.format(accuracy))

#模型损失图
plt.plot(train_loss,color='darkorange')
plt.xlabel('训练次数')
plt.ylabel('训练损失')
plt.title('LSTM模型损失函数图')
plt.show()
